import pandas as pd

from common import S_CSV_ENCODING


def _guess_usecols(filename, columns):
    df = pd.read_csv(filename, compression="zip", nrows=0)
    header = df.columns.to_list()
    usecols = [header.index(column) for column in columns]
    return usecols


def read_csv(filename, columns, dtype=None):
    usecols = _guess_usecols(filename, columns)

    df = pd.read_csv(
        filename,
        compression="zip",
        usecols=usecols,
        skiprows=1,
        header=None,
        names=columns,
        dtype="string" if dtype is None else dtype,
    )

    return df


def read_gongdan(f, usecols=None):
    df = pd.read_csv(
        f,
        usecols=usecols,
        dtype="string",
        encoding=S_CSV_ENCODING,
    )
    df.rename(
        {"派单日期": "派单时间", "结单时间": "归档时间"},
        inplace=True,
        axis=1,
    )
    df["派单时间"] = pd.to_datetime(
        df["派单时间"], format="%Y/%m/%d %H:%M"
    )
    df["归档时间"] = pd.to_datetime(
        df["归档时间"], format="%Y/%m/%d %H:%M"
    )

    return df
